# -*- coding: utf-8 -*-
import numpy as np
from sklearn import datasets

iris = datasets.load_iris()

X = iris["data"]
y = iris["target"]

# 封装为函数

def train_test_split( X, y, test_ratio=0.2, seed = None ):
    '''
    根据数据集X和y，按照test_ratio分割成X_train, X_test, y_train, y_test
    '''
    # 便于调试，设置seed
    if seed:
        np.random.seed( seed )
    
    shuffle_indexes = np.random.permutation( len( X ) )

    # 设置测试数据比例,为20%
    test_size = int(len( X )  * test_ratio) # 有可能出现浮点数，强制转化为int
    
    # 定义训练集的索引，测试集的索引
    test_indexes = shuffle_indexes[ :test_size ]
    train_indexes = shuffle_indexes[ test_size: ]
    
    
    X_test = X[test_indexes]
    X_train = X[train_indexes]
    
    y_test = y[test_indexes]
    y_train = y[train_indexes]
    
    return X_train, X_test, y_train, y_test

if __name__ == "__main__":
    from knn封装 import KNNClassifier
    X_train, X_test, y_train, y_test = train_test_sqlit( X, y )
    my_knn = KNNClassifier( k = 6 )
    my_knn.fit( X_train, y_train )
    predict = my_knn.predict( X_test )
    
    accuracy = np.sum(predict == y_test)/len( X_test ) #0.96666666666666667
    
    # sklearn 中的train_test_split
    from sklearn.model_selection import train_test_split
    X_train, X_test, y_train, y_test = train_test_split( X, y, test_size = 0.2 )
    print( X_train )
    print( X_test )
    print( y_train )
    print( y_test )

